{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5d8c4c62-eeb6-4857-a23c-08b4bbf2412f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "# 读取电影评分数据\n",
    "ratings_df = pd.read_csv(\"../data/llm-pretrain-data/ratings_title.csv\", sep=\",\", low_memory=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3b1ee724-34aa-4787-9669-2a1c0df7b5d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>userId</th>\n",
       "      <th>movieId</th>\n",
       "      <th>rating</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>tmdbId</th>\n",
       "      <th>title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>296</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147880044</td>\n",
       "      <td>680</td>\n",
       "      <td>Pulp Fiction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>306</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1147868817</td>\n",
       "      <td>110</td>\n",
       "      <td>Three Colors: Red</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>307</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147868828</td>\n",
       "      <td>108</td>\n",
       "      <td>Three Colors: Blue</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>665</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147878820</td>\n",
       "      <td>11902</td>\n",
       "      <td>Underground</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>899</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1147868510</td>\n",
       "      <td>872</td>\n",
       "      <td>Singin' in the Rain</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948967</th>\n",
       "      <td>162541</td>\n",
       "      <td>50872</td>\n",
       "      <td>4.5</td>\n",
       "      <td>1240953372</td>\n",
       "      <td>2062</td>\n",
       "      <td>Ratatouille</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948968</th>\n",
       "      <td>162541</td>\n",
       "      <td>55768</td>\n",
       "      <td>2.5</td>\n",
       "      <td>1240951998</td>\n",
       "      <td>5559</td>\n",
       "      <td>Bee Movie</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948969</th>\n",
       "      <td>162541</td>\n",
       "      <td>56176</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1240950697</td>\n",
       "      <td>6477</td>\n",
       "      <td>Alvin and the Chipmunks</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948970</th>\n",
       "      <td>162541</td>\n",
       "      <td>58559</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1240953434</td>\n",
       "      <td>155</td>\n",
       "      <td>The Dark Knight</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948971</th>\n",
       "      <td>162541</td>\n",
       "      <td>63876</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1240952515</td>\n",
       "      <td>10139</td>\n",
       "      <td>Milk</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>24948972 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          userId movieId rating   timestamp tmdbId                    title\n",
       "0              1     296    5.0  1147880044    680             Pulp Fiction\n",
       "1              1     306    3.5  1147868817    110        Three Colors: Red\n",
       "2              1     307    5.0  1147868828    108       Three Colors: Blue\n",
       "3              1     665    5.0  1147878820  11902              Underground\n",
       "4              1     899    3.5  1147868510    872      Singin' in the Rain\n",
       "...          ...     ...    ...         ...    ...                      ...\n",
       "24948967  162541   50872    4.5  1240953372   2062              Ratatouille\n",
       "24948968  162541   55768    2.5  1240951998   5559                Bee Movie\n",
       "24948969  162541   56176    2.0  1240950697   6477  Alvin and the Chipmunks\n",
       "24948970  162541   58559    4.0  1240953434    155          The Dark Knight\n",
       "24948971  162541   63876    5.0  1240952515  10139                     Milk\n",
       "\n",
       "[24948972 rows x 6 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ratings_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b0b06768-af53-4f5b-a315-cdadbec9419f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>userId</th>\n",
       "      <th>movieId</th>\n",
       "      <th>rating</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>tmdbId</th>\n",
       "      <th>title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>296</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147880044</td>\n",
       "      <td>680</td>\n",
       "      <td>Pulp Fiction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>306</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1147868817</td>\n",
       "      <td>110</td>\n",
       "      <td>Three Colors: Red</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>307</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147868828</td>\n",
       "      <td>108</td>\n",
       "      <td>Three Colors: Blue</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>665</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147878820</td>\n",
       "      <td>11902</td>\n",
       "      <td>Underground</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>899</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1147868510</td>\n",
       "      <td>872</td>\n",
       "      <td>Singin' in the Rain</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948967</th>\n",
       "      <td>162541</td>\n",
       "      <td>50872</td>\n",
       "      <td>4.5</td>\n",
       "      <td>1240953372</td>\n",
       "      <td>2062</td>\n",
       "      <td>Ratatouille</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948968</th>\n",
       "      <td>162541</td>\n",
       "      <td>55768</td>\n",
       "      <td>2.5</td>\n",
       "      <td>1240951998</td>\n",
       "      <td>5559</td>\n",
       "      <td>Bee Movie</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948969</th>\n",
       "      <td>162541</td>\n",
       "      <td>56176</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1240950697</td>\n",
       "      <td>6477</td>\n",
       "      <td>Alvin and the Chipmunks</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948970</th>\n",
       "      <td>162541</td>\n",
       "      <td>58559</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1240953434</td>\n",
       "      <td>155</td>\n",
       "      <td>The Dark Knight</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24948971</th>\n",
       "      <td>162541</td>\n",
       "      <td>63876</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1240952515</td>\n",
       "      <td>10139</td>\n",
       "      <td>Milk</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>24948972 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          userId movieId rating   timestamp tmdbId                    title\n",
       "0              1     296    5.0  1147880044    680             Pulp Fiction\n",
       "1              1     306    3.5  1147868817    110        Three Colors: Red\n",
       "2              1     307    5.0  1147868828    108       Three Colors: Blue\n",
       "3              1     665    5.0  1147878820  11902              Underground\n",
       "4              1     899    3.5  1147868510    872      Singin' in the Rain\n",
       "...          ...     ...    ...         ...    ...                      ...\n",
       "24948967  162541   50872    4.5  1240953372   2062              Ratatouille\n",
       "24948968  162541   55768    2.5  1240951998   5559                Bee Movie\n",
       "24948969  162541   56176    2.0  1240950697   6477  Alvin and the Chipmunks\n",
       "24948970  162541   58559    4.0  1240953434    155          The Dark Knight\n",
       "24948971  162541   63876    5.0  1240952515  10139                     Milk\n",
       "\n",
       "[24948972 rows x 6 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 确保每个用户至少有10部电影评分\n",
    "min_unique_movies_per_user = 10\n",
    "enough_movies_df = ratings_df.groupby('userId').filter(lambda x: len(x) >= 10)\n",
    "enough_movies_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "df594d79-2121-4f2f-bbed-373a7b443114",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11657558     75732\n",
       "16528004    107465\n",
       "18190901    118047\n",
       "11162538     72533\n",
       "12518123     81115\n",
       "             ...  \n",
       "19044799    123757\n",
       "2633499      17549\n",
       "854908        5767\n",
       "12909539     83754\n",
       "20561506    134012\n",
       "Name: userId, Length: 1000, dtype: object"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 随机数种子：保证每次重复运行程序时抽取时，样本保持不变\n",
    "random_state = 6\n",
    "# 设定抽取 1000 用户\n",
    "selected_users = enough_movies_df['userId'].drop_duplicates().sample(n=1000, random_state=random_state)\n",
    "selected_users"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b016c7f6-9d23-439e-83bd-3c57a0de6385",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>userId</th>\n",
       "      <th>movieId</th>\n",
       "      <th>rating</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>tmdbId</th>\n",
       "      <th>title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>15414968</th>\n",
       "      <td>100096</td>\n",
       "      <td>809</td>\n",
       "      <td>3.0</td>\n",
       "      <td>857857780</td>\n",
       "      <td>18550</td>\n",
       "      <td>Fled</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15414957</th>\n",
       "      <td>100096</td>\n",
       "      <td>260</td>\n",
       "      <td>5.0</td>\n",
       "      <td>857857090</td>\n",
       "      <td>11</td>\n",
       "      <td>Star Wars</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15414953</th>\n",
       "      <td>100096</td>\n",
       "      <td>100</td>\n",
       "      <td>4.0</td>\n",
       "      <td>857857161</td>\n",
       "      <td>11062</td>\n",
       "      <td>City Hall</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15414954</th>\n",
       "      <td>100096</td>\n",
       "      <td>102</td>\n",
       "      <td>3.0</td>\n",
       "      <td>857857232</td>\n",
       "      <td>47475</td>\n",
       "      <td>Mr. Wrong</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15414959</th>\n",
       "      <td>100096</td>\n",
       "      <td>648</td>\n",
       "      <td>3.0</td>\n",
       "      <td>857856991</td>\n",
       "      <td>954</td>\n",
       "      <td>Mission: Impossible</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1492041</th>\n",
       "      <td>9991</td>\n",
       "      <td>2067</td>\n",
       "      <td>4.0</td>\n",
       "      <td>939406876</td>\n",
       "      <td>907</td>\n",
       "      <td>Doctor Zhivago</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1492040</th>\n",
       "      <td>9991</td>\n",
       "      <td>1961</td>\n",
       "      <td>2.0</td>\n",
       "      <td>939406876</td>\n",
       "      <td>380</td>\n",
       "      <td>Rain Man</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1492023</th>\n",
       "      <td>9991</td>\n",
       "      <td>1233</td>\n",
       "      <td>5.0</td>\n",
       "      <td>939406957</td>\n",
       "      <td>387</td>\n",
       "      <td>Das Boot</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1492010</th>\n",
       "      <td>9991</td>\n",
       "      <td>750</td>\n",
       "      <td>4.0</td>\n",
       "      <td>939405822</td>\n",
       "      <td>935</td>\n",
       "      <td>Dr. Strangelove or: How I Learned to Stop Worr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1492049</th>\n",
       "      <td>9991</td>\n",
       "      <td>2657</td>\n",
       "      <td>2.0</td>\n",
       "      <td>939406129</td>\n",
       "      <td>36685</td>\n",
       "      <td>The Rocky Horror Picture Show</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10000 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          userId movieId rating  timestamp tmdbId  \\\n",
       "15414968  100096     809    3.0  857857780  18550   \n",
       "15414957  100096     260    5.0  857857090     11   \n",
       "15414953  100096     100    4.0  857857161  11062   \n",
       "15414954  100096     102    3.0  857857232  47475   \n",
       "15414959  100096     648    3.0  857856991    954   \n",
       "...          ...     ...    ...        ...    ...   \n",
       "1492041     9991    2067    4.0  939406876    907   \n",
       "1492040     9991    1961    2.0  939406876    380   \n",
       "1492023     9991    1233    5.0  939406957    387   \n",
       "1492010     9991     750    4.0  939405822    935   \n",
       "1492049     9991    2657    2.0  939406129  36685   \n",
       "\n",
       "                                                      title  \n",
       "15414968                                               Fled  \n",
       "15414957                                          Star Wars  \n",
       "15414953                                          City Hall  \n",
       "15414954                                          Mr. Wrong  \n",
       "15414959                                Mission: Impossible  \n",
       "...                                                     ...  \n",
       "1492041                                      Doctor Zhivago  \n",
       "1492040                                            Rain Man  \n",
       "1492023                                            Das Boot  \n",
       "1492010   Dr. Strangelove or: How I Learned to Stop Worr...  \n",
       "1492049                       The Rocky Horror Picture Show  \n",
       "\n",
       "[10000 rows x 6 columns]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# 对每个被抽取的用户，找到他们的最新10条电影评分\n",
    "latest_user_ratings = (\n",
    "    enough_movies_df[enough_movies_df['userId'].isin(selected_users)]\n",
    "    .sort_values(by=['userId', 'timestamp'], ascending=[True, False])  # 按照用户ID升序，评分降序排列\n",
    "    .groupby('userId')\n",
    "    .sample(n=10, random_state=random_state)\n",
    ")\n",
    "latest_user_ratings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ccb38aa3-c4f7-450b-8b28-ae210d78c0a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 确保每个被选中的用户确实得到了他们最新的10条电影评分，否则程序停止\n",
    "latest_user_ratings.groupby('userId').size().min() == 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64cc7e3-3fe3-4175-b34a-c9c253dba41a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# lowest_ratings = (\n",
    "#     enough_movies_df[enough_movies_df['userId'].isin(selected_users)]\n",
    "#     .sort_values(by=['userId', 'rating'], ascending=[True, True])  # 按照用户ID升序，评分升序排列\n",
    "#     .groupby('userId')\n",
    "#     .head(5)\n",
    "# )\n",
    "\n",
    "# 合并高低评分\n",
    "# latest_user_ratings = pd.concat([highest_ratings, lowest_ratings]).sort_values(by=['userId', 'rating'], ascending=[True, False])\n",
    "\n",
    "\n",
    "# 对每个用户最新的10条评分进行切片，前9条作为训练集，后1条作为测试集\n",
    "user_set = latest_user_ratings.sort_values(by=['userId', 'timestamp'], ascending=[True, False])\n",
    "test_set = user_set.groupby('userId').head(1)\n",
    "train_set = user_set.drop(test_set.index)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e22c56f2-6763-4d84-9ed3-0cbab55fc633",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 确保每个用户在训练集和测试集中都有数据\n",
    "# assert train_set.groupby('userId').size().min() == train_test_split_count\n",
    "# assert test_set.groupby('userId').size().min() == 10 - train_test_split_count\n",
    "\n",
    "# 输出处理后的数据集\n",
    "train_set.to_csv(f\"../data/llm-pretrain-data/train_set_{random_state}_1k.csv\", index=False)\n",
    "test_set.to_csv(f\"../data/llm-pretrain-data/test_set_{random_state}_1k.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffc14cb-c6e3-4483-950e-56948e3692a4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
